#![allow(async_fn_in_trait)]

use content::{ContentCollection, ContentStore};
use serde::{Deserialize, Serialize};

pub mod content;
pub mod engine;

pub mod train;

//pub mod mixed;

#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct SearchConfig {
    ///结果数目
    pub max_result_length: Option<usize>,
    ///是否允许空结果
    pub allow_null_result: bool,
    ///设置可视为正常结果的最大长度
    pub max_distance: Option<f32>,
}

impl Default for SearchConfig {
    fn default() -> Self {
        Self {
            max_result_length: Some(3),
            allow_null_result: true,
            max_distance: None,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchResultItem {
    id: usize,
    pub distance: f32,
}

impl PartialEq for SearchResultItem {
    fn eq(&self, other: &Self) -> bool {
        self.distance.eq(&other.distance) && self.id.eq(&other.id)
    }
}
impl Eq for SearchResultItem {}
impl Ord for SearchResultItem {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        match self.distance.partial_cmp(&other.distance) {
            Some(s) => s,
            None => self.id.cmp(&other.id),
        }
    }
}
impl PartialOrd for SearchResultItem {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

#[cfg(test)]
mod prepare {
    use std::fs::File;

    use crate::{
        content::Content,
        engine::{simple, Engine, VecEngineImpl},
    };

    use self::{
        content::ContentInnerUnembeded, engine::simple::SimpleEngine, train::TrainableEngine,
    };

    use super::*;
    use anyhow::Result;
    use embd::Embd;
    use embedding_model::EmbeddingModel;

    const UNEMBEDED_PATH: &str = "./materials/blank_unembeded1.json";
    const EMBEDED_PATH: &str = "./materials/contentstore.json";
    const SIMPLE_ENGINE_INDEX_PATH: &str = "./materials/se_index.json";
    #[cfg(target_family = "windows")]
    const EMBEDDING_MODEL_PATH: &str = r"..\embd\model\text2vec-base-chinese_safetensors\";
    #[cfg(target_family = "unix")]
    const EMBEDDING_MODEL_PATH: &str = r"../embd/model/text2vec-base-chinese_safetensors/";

    #[test]
    fn store_blank_unembeded() -> Result<()> {
        let wtr = File::create(UNEMBEDED_PATH)?;
        let val = ContentInnerUnembeded::blank();
        serde_json::to_writer(wtr, &vec![val])?;
        Ok(())
    }

    #[test]
    fn convert_unembeded_toembeded() -> Result<()> {
        let model = Embd::load(EMBEDDING_MODEL_PATH, embd::DType::F32, &embd::Device::Cpu)?;
        let unembeded = ContentInnerUnembeded::from_json_file(UNEMBEDED_PATH)?;
        let embeded = ContentStore::from_unembeded(unembeded, &model, 1)?;
        embeded.save_json_file(EMBEDED_PATH)?;
        Ok(())
    }

    #[test]
    fn build_simple_engine_index() -> Result<()> {
        let config = engine::simple::InitConfig {
            path: SIMPLE_ENGINE_INDEX_PATH.into(),
        };
        let mut engine = SimpleEngine::blank(&config)?;
        let material = ContentStore::from_json_file(EMBEDED_PATH)?;
        engine.train(material)?;
        engine.save(&config)?;
        Ok(())
    }

    #[tokio::test]
    async fn test_simple_engine_search() -> Result<()> {
        let config = engine::simple::InitConfig {
            path: SIMPLE_ENGINE_INDEX_PATH.into(),
        };
        let engine =
            VecEngineImpl::<SimpleEngine, Embd>::new::<SimpleEngine, simple::InitConfig, &str>(
                &config,
                EMBEDDING_MODEL_PATH,
                embedding_model::DType::F32,
                &embedding_model::Device::Cpu,
            )?;
        let contents = ContentStore::from_json_file(EMBEDED_PATH)?;
        let sentence = "我购买的计算机配件怎么报销";
        let search_config = SearchConfig {
            max_result_length: Some(2),
            allow_null_result: true,
            max_distance: Some(600f32),
        };
        let res = engine
            .search(&sentence, &contents, &search_config)
            .await
            .unwrap();
        println!("{} items found.", res.len());
        for (index, content) in res.into_iter().enumerate() {
            println!("{index} => {}", content.content.content());
        }
        Ok(())
    }
}
